#include <iostream>
#include <NvOnnxParser.h>
#include "../include/misc.h"

using namespace std;

void print_ioinfo(nvinfer1::INetworkDefinition* network){
    int input_nums = network->getNbInputs();
    for(int i=0;i<input_nums;i++){
        nvinfer1::ITensor* input = network->getInput(i);
        string input_name = input->getName();
        nvinfer1::Dims input_dims = input->getDimensions();
        printf("input<%d> tensor name = %s, dim = ", i,input_name.c_str());
        for(int i=0;i<input_dims.nbDims;i++){
            cout<<input_dims.d[i]<<"  ";
        }
        cout<<endl;
    }
    int output_nums = network->getNbOutputs();
    for(int i=0;i<output_nums;i++){
        nvinfer1::ITensor* output = network->getOutput(i);
        string output_name = output->getName();
        nvinfer1::Dims output_dims = output->getDimensions();
        printf("output<%d> tensor name = %s,  dim = ", i,output_name.c_str());
        for(int i=0;i<output_dims.nbDims;i++){
            cout<<output_dims.d[i]<<"  ";
        }
        cout<<endl;
    }
}

int main(){
    string root_path = "/files/code/Csdn-Source-Code-Implementation/tensorrt/step0-diymodel/weights/";
    string onnx_file = root_path+"best.onnx";
    string trt_file = root_path+"best.trt";

    TRTLogger logger;
    nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
    nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
    nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1);
    nvonnxparser::IParser* parser = nvonnxparser::createParser(*network,logger);

    if(!parser->parseFromFile(onnx_file.c_str(),1)){
        printf("Failed to parse classifier.onnx\n");
        // 注意这里的几个指针还没有释放，是有内存泄漏的，后面考虑更优雅的解决
        return -1;
    }
    printf("succeed parse onnx\n");

    config->setMaxWorkspaceSize(1<<28);
    printf("WorkspaceSize = %.2f MB\n",(1<<28)/1024.0/1024.0);
    nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile();
    nvinfer1::ITensor* input = network->getInput(0);
    nvinfer1::Dims  input_dim = input->getDimensions();
    input_dim.d[0] = 1;
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMIN, input_dim);
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kOPT, input_dim);
    input_dim.d[0] = 10;
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, input_dim);
    config->addOptimizationProfile(profile);

    print_ioinfo(network);

    nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network,*config);
    if(engine== nullptr){
        printf("Build engine failed.\n");
        return -1;
    }

    nvinfer1::IHostMemory* model_data = engine->serialize();

    FILE *f = fopen(trt_file.c_str(), "wb");
    fwrite(model_data->data(), 1, model_data->size(), f);
    fclose(f);

    model_data->destroy();
    parser->destroy();
    engine->destroy();
    network->destroy();
    config->destroy();
    builder->destroy();
    printf("done\n");
    return 0;
}